#pragma once

#include "RayTracingLibrary.h"

class RayTracingPipeline : public RHIPipeline
{
	struct BuilderDetails {

		RefCountPtr<RHIDescriptor> descriptor;

		std::vector<ShaderPtr<RHIShader>> shaderStages;
		VkRayTracingShaderGroupCreateInfoKHR tempCroup;
		std::vector<VkRayTracingShaderGroupCreateInfoKHR> shaderGroups;
		uint8* dataPtr;
		uint64 dataSize;
		std::array<std::unordered_map<uint32, std::vector<uint8>>, 4> dataGroups;

		uint32 maxPipelineRayRecursionDepth;

		std::vector<RefCountPtr<RayTracingLibrary>> rayTracingLibrarys;

		std::vector<VkDynamicState> dynamicStates;

		RefCountPtr<RHIPipeline> basePipelineHandle;
		int32 basePipelineIndex;

		BuilderDetails()
		{
			tempCroup = {};
			tempCroup.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR;
			tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
			tempCroup.generalShader = VK_SHADER_UNUSED_KHR;
			tempCroup.closestHitShader = VK_SHADER_UNUSED_KHR;
			tempCroup.anyHitShader = VK_SHADER_UNUSED_KHR;
			tempCroup.intersectionShader = VK_SHADER_UNUSED_KHR;

			dataPtr = nullptr;
			dataSize = 0;

			maxPipelineRayRecursionDepth = 0;

			// 父管线
			basePipelineIndex = -1;
		}
	};

public:

	class Builder : public BuilderBase<BuilderDetails> {
		friend struct BuilderDetails;
	public:
		Builder() noexcept;
		Builder(Builder const& rhs) noexcept;
		Builder(Builder&& rhs) noexcept;
		~Builder() noexcept;
		Builder& operator=(Builder const& rhs) noexcept;
		Builder& operator=(Builder&& rhs) noexcept;

		/************************************************************************/
		/*                              Descriptor                              */
		/************************************************************************/

		Builder& SetDescriptor(RefCountPtr<RHIDescriptor> descriptor) noexcept;

		/************************************************************************/
		/*                           ShaderStage                                */
		/************************************************************************/

		Builder& AddShaderStage(ShaderPtr<RHIShader> shader) noexcept;

		/************************************************************************/
		/*                           ShaderGroup                                */
		/************************************************************************/

		Builder& PushShader(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushRaygen(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushMiss(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushClosestHit(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushAnyHit(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushIntersection(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushCallable(ShaderPtr<RHIShader> shader) noexcept;
		template<typename T>
		Builder& PushData(T& data) noexcept
		{
			PushData((uint8*)&data, sizeof(T));
			return *this;
		}
		Builder& PackageGroup() noexcept;

		/************************************************************************/
		/*                           Ray   Param                                */
		/************************************************************************/

		Builder& SetMaxPipelineRayRecursionDepth(uint32 maxPipelineRayRecursionDepth) noexcept;

		/************************************************************************/
		/*                           Ray   Param                                */
		/************************************************************************/
		Builder& AddRayTracingLibrary(RefCountPtr<RayTracingLibrary> rayTracingLibrary) noexcept;

		/************************************************************************/
		/*                           DynamicState                               */
		/************************************************************************/

		Builder& AddDynamicState(VkDynamicState dynamicState) noexcept;

		/************************************************************************/
		/*                           BasePipeline                               */
		/************************************************************************/

		Builder& SetBasePipeline(RefCountPtr<RHIPipeline> basePipeline, int32 index = -1) noexcept;


		RefCountPtr<RayTracingPipeline> Build();
		RefCountPtr<RayTracingPipeline> Build(RenderProcess* process, const std::string& name);

	protected:

		uint32 GetShaderIndex(ShaderPtr<RHIShader> shader) const;

		void PushData(uint8* data, uint64 dataSize);

	private:

		friend class RayTracingPipeline;
	};

	RayTracingPipeline(BuilderDetails& builderDetails);

	virtual ~RayTracingPipeline();

	virtual void ReleaseResource() override;

	std::vector<ShaderPtr<RHIShader>>& GetShaderStages();

	std::array<VkStridedDeviceAddressRegionKHR, 4> GetRegions() const;

	const VkStridedDeviceAddressRegionKHR& GetRaygenRegion() const;

	const VkStridedDeviceAddressRegionKHR& GetMissRegion() const;

	const VkStridedDeviceAddressRegionKHR& GetHitRegion() const;

	const VkStridedDeviceAddressRegionKHR& GetCallableRegion() const;

protected:

	void CreateShaderBingingTable();

	uint8* CreateStagingBuffer(VkDeviceSize size, VkBuffer& buffer, VkDeviceMemory& bufferMemory);

	VkDeviceAddress CreateGroupBuffer(VkDeviceSize size, VkBuffer& buffer, VkDeviceMemory& bufferMemory);

	void CopyBuffer(VkBuffer srcBuffer, VkBuffer dstBuffer, VkDeviceSize size);

protected:

	std::vector<VkRayTracingShaderGroupCreateInfoKHR> mShaderGroups;

	std::array<std::unordered_map<uint32, std::vector<uint8>>, 4> mDataGroups;

	uint32 mMaxPipelineRayRecursionDepth;

	std::vector<RefCountPtr<RayTracingLibrary>> mRayTracingLibrarys;

	std::array<VkBuffer, 4> mGroupBuffers;

	std::array<VkDeviceMemory, 4> mGroupMemorys;

	std::array<VkStridedDeviceAddressRegionKHR, 4> mGroupRegions;

	std::vector<ShaderPtr<RHIShader>> mShaderStages;
};